
import pymupdf as fitz
from PIL import Image
from .handle_pdf.fitz import page_to_image
from .handle_img.tesseract import ocr_image
from .invoices.normal import get_normal_invoice_info_from_page
from .invoices.railway import get_railway_invoice_info_from_page
from .invoices.types_config import INVOICE_TYPES

def get_img_pattern(img):
    print(ocr_image(img)) 
    return ocr_image(img)
    pass

def get_invoice_type(page_text)-> str | None:
    for invoice_type, origin_text in INVOICE_TYPES.items():
        if origin_text in page_text:
            return invoice_type
    return None

def get_file_type(file_path):
    if "." in file_path:
        file_type=file_path.split('.')[-1].lower()
    else:
        file_type=file_path.split('/')[-1].lower()
    return file_type

def parse_invoice(file_path):
    file_type=get_file_type(file_path)
    match file_type:
        case "pdf":
            # 处理PDF（优先提取文本，若失败则OCR识别图片）
            doc = fitz.open(file_path)
            for page in doc:
                # 尝试提取原生文本
                page_text = page.get_text("text", sort=True)
                if len(page_text.strip()) > 10:  # 简单判断是否有有效文本
                    #TODO 提取发票类型,并提交相关模块进行处理
                    invoice_type = get_invoice_type(page_text)
                    match  invoice_type:
                        case "railway":
                            result = get_railway_invoice_info_from_page(page)
                        case "normal":
                            result = get_normal_invoice_info_from_page(page)
                else:
                    # 将 PDF 页面转为图像进行 OCR
                    img = page_to_image(page, dpi=300)
                    result= get_img_pattern(img)
            doc.close()
            return result
        case "jpg"|"png":
            # 处理图片（JPG/PNG）
            img = Image.open(file_path)
            return get_img_pattern(img)

